package com.jboltai.capability.vdb.types.postgresql;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import com.jboltai.capability.embedding.JBoltEmbedding;
import com.jboltai.capability.vdb.EmbeddingException;
import com.jboltai.capability.vdb.JBoltVDB;
import com.jboltai.capability.vdb.VDBOptResult;
import com.jboltai.capability.vdb.collection.VDBCollection;
import com.jboltai.capability.vdb.collection.VDBMetricType;
import com.jboltai.capability.vdb.collection.field.VDBField;
import com.jboltai.capability.vdb.collection.field.VDBFieldType;
import com.jboltai.capability.vdb.collection.field.VectorIndexType;
import com.jboltai.capability.vdb.doc.VDBDoc;
import com.jboltai.capability.vdb.search.VDBConditions;
import com.jboltai.capability.vdb.types.tencent.FieldUtil;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.embedding.EmbeddingEvent;
import com.jboltai.resource.embedding.EmbeddingModel;
import com.jboltai.resource.vdb.VDBResource;
import org.postgresql.jdbc.PgResultSetMetaData;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static com.jboltai.capability.vdb.collection.field.VDBFieldType.VECTOR;

public class JBoltPostgreSQLVDB extends JBoltVDB {
    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    public JBoltPostgreSQLVDB(VDBResource resource) {
        super(resource);
        if (StrUtil.isBlank(resource.getDatabase())) {
            throw new RuntimeException("向量数据库资源[" + resource.getId() + "]未设置database");
        }
        if (EmbeddingModel.isTencentInnerModel(resource.getEmbeddingModel())) {
            throw new RuntimeException(StrUtil.format("向量数据库资源[{}]不支持使用[{}]Embedding模型", resource.getId(), resource.getEmbeddingModel()));
        }
    }

    @Override
    public VDBOptResult<Boolean> createCollection(VDBCollection collection) {
        if (resource.getEmbeddingModel() == null) {
            throw new RuntimeException("向量数据库[" + resource.getId() + "]未指定Embedding模型");
        }
        Connection conn = null;
        try {
            String createTableSQL = createTableSQL(collection);
            String createIndexSQL = createIndexSQL(collection);


            conn = getConnFromResource();
            Statement stmt = conn.createStatement();
            Integer result = stmt.executeUpdate(createTableSQL);
            if (result != 0) {
                return failData(false);
            }
            result = stmt.executeUpdate(createIndexSQL);
            return result == 0 ? successData(true) : failData(false);
        } catch (Exception e) {
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            if (conn != null) {
                try {
                    conn.close();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }

    private String createTableSQL(VDBCollection collection) {
        StringBuilder fieldSB = new StringBuilder();
        for (VDBField field : collection.getFields()) {
            VDBFieldType vdbFieldType = field.getVdbFieldType();
            switch (vdbFieldType) {
                case STRING: {
                    fieldSB.append(field.getName()).append(" varchar(").append(field.isPrimaryKey() ? 40 : 65535).append(")");
                    break;

                }
                case NUMBER32:
                    fieldSB.append(field.getName()).append(" int4");
                    break;
                case NUMBER64:
                    fieldSB.append(field.getName()).append(" int8");
                    break;
                case VECTOR:
                    fieldSB.append(field.getName()).append(" vector(").append(resource.getEmbeddingModel().getDimension()).append(")");
                    break;
                default:
                    throw new RuntimeException("未提供的支持");
            }
            if (field.isPrimaryKey()) {
                fieldSB.append(" PRIMARY KEY");
            }
            fieldSB.append(",");
        }
        String template = "CREATE TABLE %s (%s)";
        String sql = String.format(template, collection.getName(), fieldSB.substring(0, fieldSB.length() - 1));
        return sql;
    }

    private String createIndexSQL(VDBCollection collection) {
        VDBField vectorField = null;
        // 找到向量字段
        for (VDBField field : collection.getFields()) {
            VDBFieldType vdbFieldType = field.getVdbFieldType();
            if (vdbFieldType == VECTOR) {
                vectorField = field;
            }
        }
        if (vectorField == null) {
            throw new RuntimeException("未设置向量字段");
        }

        String tableName = collection.getName();
        VectorIndexType vectorIndexType = vectorField.getVectorIndexType();
        VDBMetricType vdbMetricType = vectorField.getVdbMetricType();
        if (vectorIndexType == VectorIndexType.HNSW) {
            JSONObject param = vectorField.getVectorIndexParams();
            Integer M = param.getInteger("M");
            Integer efConstruction = param.getInteger("efConstruction");
            switch (vdbMetricType) {
                case L2:
                    return "CREATE INDEX ON " + tableName + " USING hnsw (" + vectorField.getName() + " vector_l2_ops) WITH (m = " + M + ", ef_construction = " + efConstruction + ")";
                case IP:
                    return "CREATE INDEX ON " + tableName + " USING hnsw (" + vectorField.getName() + " vector_ip_ops) WITH (m = " + M + ", ef_construction = " + efConstruction + ")";
                case DEFAULT:
                    return "CREATE INDEX ON " + tableName + " USING hnsw (" + vectorField.getName() + " vector_cosine_ops) WITH (m = " + M + ", ef_construction = " + efConstruction + ")";
            }
            throw new RuntimeException("不支持的VDBMetricType");
        } else if (vectorIndexType == VectorIndexType.IVF_FLAT) {
            switch (vdbMetricType) {
                case L2:
                    return "CREATE INDEX ON " + tableName + " USING ivfflat (" + vectorField.getName() + " vector_l2_ops) WITH (lists = 100)";
                case IP:
                    return "CREATE INDEX ON " + tableName + " USING ivfflat (" + vectorField.getName() + " vector_ip_ops) WITH (lists = 100)";
                case DEFAULT:
                    return "CREATE INDEX ON " + tableName + " USING ivfflat (" + vectorField.getName() + " vector_cosine_ops) WITH (lists = 100)";
            }
            throw new RuntimeException("不支持的VDBMetricType");
        } else {
            throw new RuntimeException("不支持的向量索引类型");
        }
    }

    @Override
    public VDBOptResult<Boolean> dropCollection(String tableName) {
        String template = "DROP TABLE IF EXISTS %s";
        String sql = String.format(template, tableName);
        Connection conn = null;
        try {
            conn = getConnFromResource();
            Boolean result = executeUpdate(conn, sql);
            return resultData(result, result);
        } catch (Exception e) {
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            closeConnection(conn);
        }
    }

    @Override
    public VDBOptResult<Boolean> checkCollectionExist(String tableName) {
        String sql = "SELECT COUNT(*) AS count FROM information_schema.tables WHERE table_name = '%s'";
        sql = String.format(sql, tableName);

        Connection conn = getConnFromResource();
        try {
            return executeCount(conn, sql) == 1 ? successData(true) : failData(false);
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            closeConnection(conn);
        }
    }

    private Connection getConnFromResource() {
        PostgreSQLVDBResource resource = (PostgreSQLVDBResource) getResource();
        return resource.getConnection();
    }

    @Override
    public VDBOptResult<Boolean> truncateCollection(String tableName) {
        String sql = "TRUNCATE TABLE %s";
        sql = String.format(sql, tableName);

        Connection conn = getConnFromResource();
        try {
            return executeUpdate(conn, sql) ? successData(true) : failData(false);
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            closeConnection(conn);
        }
    }

    @Override
    public VDBOptResult<Integer> insert(String tableName, VDBDoc... docs) {
        Connection conn = null;
        try {
            String sql = createInsertSQL(tableName, docs);
            conn = getConnFromResource();
            Statement stmt = conn.createStatement();
            Integer count = stmt.executeUpdate(sql);
            return successData(count);
        } catch (Exception e) {
            if (e instanceof EmbeddingException) {
                return new VDBOptResult<>("Embedding发生异常", e);
            }
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            closeConnection(conn);
        }
    }

    private String createInsertSQL(String tableName, VDBDoc... docs) {
        String template = "INSERT INTO %s (%s) VALUES (%s);";
        StringBuilder sql = new StringBuilder();

        for (VDBDoc doc : docs) {
            StringBuilder columnSB = new StringBuilder();
            StringBuilder columnValueSB = new StringBuilder();

            if (StrUtil.isNotBlank(doc.getId())) {
                columnSB.append("id");
                columnValueSB.append("'").append(doc.getId()).append("'");
            }

            if (StrUtil.isNotBlank(doc.getDoc())) {
                String docStr = doc.getDoc();
                String vector = embeddingDoc(docStr);

                columnSB.append(",doc");
                columnValueSB.append(",'").append(docStr).append("'");

                columnSB.append(",vector");
                columnValueSB.append(",'").append(vector).append("'");
            }

            Map<String, Object> fieldMap = doc.getFieldValues();
            for (String key : fieldMap.keySet()) {
                columnSB.append(",").append(key);
                Object value = fieldMap.get(key);
                String clazzType = value.getClass().getSimpleName();
                if (FieldUtil.isNumber(clazzType)) {
                    columnValueSB.append(",").append(fieldMap.get(key));
                } else {
                    columnValueSB.append(",").append("'").append(fieldMap.get(key)).append("'");
                }
            }
            sql.append(String.format(template, tableName, columnSB, columnValueSB));
        }
        return sql.toString();
    }

    private String embeddingDoc(String doc) {
        EmbeddingEvent event = JBoltEmbedding.embedding(resource.getEmbeddingModel(), doc);
        event.onSuccess((e, r) -> {
            event.setAttr("result", r.get(0));
        }).onFail((e, err) -> {
            Throwable t = err.getException();
            throw new RuntimeException(t);
        });
        event.publish().await();
        if (event.isFail()) {
            throw new EmbeddingException(event.getError());
        }
        return event.getAttr("result").toString();
    }

    @Override
    public VDBOptResult<List<VDBDoc>> search(String tableName, VDBConditions conditions) {
        if (tableName == null) {
            throw new RuntimeException("未指定tableName");
        }

        String vectorValue = "";

        List<String> matchTextList = conditions.getMatchText();
        if (matchTextList != null && matchTextList.size() > 1) {
            throw new RuntimeException("仅支持单条向量检索");
        } else if (matchTextList != null && matchTextList.size() == 1) {
            vectorValue = embeddingDoc(matchTextList.get(0));
        }

        String sql = SQLBuilder.buildSelectSQL(tableName, vectorValue, conditions);

        List<VDBDoc> docList = new ArrayList<>();

        Connection conn = null;
        try {
            conn = getConnFromResource();
            Statement stmt = conn.createStatement();
            ResultSet rs = stmt.executeQuery(sql);
            while (rs.next()) {
                PgResultSetMetaData metaData = (PgResultSetMetaData) rs.getMetaData();
                Integer columnCount = metaData.getColumnCount();
                VDBDoc doc = new VDBDoc();
                for (int i = 1; i <= columnCount; i++) {
                    String columnName = metaData.getColumnName(i);
                    switch (columnName) {
                        case VDBField.ID_FIELD_NAME:
                            doc.setId(rs.getString(i));
                            continue;
                        case VDBField.DOC_FIELD_NAME:
                            doc.setDoc(rs.getString(i));
                            continue;
                        case VDBField.VECTOR_FIELD_NAME:
                            String vector = rs.getString(i);
                            doc.setVector(strToDoubleArr(vector));
                            continue;
                        case "score":
                            doc.setScore(rs.getDouble(i));
                            continue;
                        default:
                            doc.setField(columnName, rs.getString(i));
                            continue;
                    }
                }
                docList.add(doc);
            }
            return successData(docList);
        } catch (Exception e) {
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            closeConnection(conn);
        }
    }

    private List<Double> strToDoubleArr(String vector) {
        vector = vector.substring(1, vector.length() - 1);
        return Arrays.stream(vector.split(",")).map(Double::parseDouble).collect(Collectors.toList());
    }

    @Override
    public VDBOptResult<Boolean> delete(String tableName, VDBConditions conditions) {
        if (tableName == null) {
            throw new RuntimeException("未指定tableName");
        }

        String sql = SQLBuilder.buildDeleteSQL(tableName, conditions);

        Connection conn = null;
        try {
            conn = getConnFromResource();
            Statement stmt = conn.createStatement();
            Integer count = stmt.executeUpdate(sql);
            Boolean result = count >= 0;
            return resultData(result, result);
        } catch (Exception e) {
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            closeConnection(conn);
        }
    }

    @Override
    public VDBOptResult<Integer> update(String tableName, VDBDoc doc, VDBConditions conditions) {
        if (tableName == null) {
            throw new RuntimeException("未指定tableName");
        }

        if (doc != null && doc.getId() != null) {
            conditions.eq(VDBField.ID_FIELD_NAME, doc.getId());
            doc.setId(null);
        }

        String vectorValue = "";
        if (doc != null && StrUtil.isNotBlank(doc.getDoc())) {
            vectorValue = embeddingDoc(doc.getDoc());
        }
        String sql = SQLBuilder.buildUpdateSQL(tableName, doc, vectorValue, conditions);
        System.out.println(sql);

        Connection conn = null;
        try {
            conn = getConnFromResource();
            Statement stmt = conn.createStatement();
            Integer count = stmt.executeUpdate(sql);
            Boolean result = count >= 0;
            return resultData(result, count);
        } catch (Exception e) {
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            closeConnection(conn);
        }
    }

    @Override
    public VDBOptResult<List<VDBCollection>> listCollections() {
        String sql = "SELECT tablename FROM pg_tables WHERE schemaname = 'public'";

        List<VDBCollection> tableList = new ArrayList<>();

        Connection conn = null;
        try {
            conn = getConnFromResource();
            Statement stmt = conn.createStatement();
            ResultSet rs = stmt.executeQuery(sql);
            while (rs.next()) {
                String tableName = rs.getString("tablename");

                VDBCollection table = new VDBCollection(this);
                table.setName(tableName);
                tableList.add(table);
            }
            return successData(tableList);
        } catch (Exception e) {
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            closeConnection(conn);
        }
    }

    @Override
    public VDBOptResult<Boolean> checkDatabaseExist(String database) {
        String template = "SELECT COUNT(*) AS count FROM pg_database WHERE datname = ?";
        Connection conn = null;
        try {
            conn = getConnFromResource();
            PreparedStatement ps = conn.prepareStatement(template);
            ps.setObject(1, database);
            ResultSet rs = ps.executeQuery();
            if (rs.next()) {
                int count = rs.getInt("count");
                boolean result = count >= 1;
                return resultData(result, result);
            }
            return failData(false);
        } catch (Exception e) {
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            closeConnection(conn);
        }
    }

    @Override
    public VDBOptResult<Void> createDatabase(String database) {
        String sql = "CREATE DATABASE " + database;

        Connection conn = getConnFromResource();
        try {
            conn.createStatement().execute(sql);
            return success();
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            closeConnection(conn);
        }
    }

    private void closeConnection(Connection conn) {
        if (conn != null) {
            try {
                conn.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    private boolean executeUpdate(Connection conn, String sql) {
        try {
            Statement stmt = conn.createStatement();
            Integer result = stmt.executeUpdate(sql);
            return result == 0;
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            if (conn != null) {
                try {
                    conn.close();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }

    private Integer executeCount(Connection conn, String sql) {
        try {
            Statement stmt = conn.createStatement();
            ResultSet rs = stmt.executeQuery(sql);
            if (rs.next()) {
                return rs.getInt("count");
            }
            return 0;
        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            closeConnection(conn);
        }
    }

    private VDBOptResult<Void> success() {
        return new VDBOptResult<>(true, null);
    }

    private <T> VDBOptResult<T> successData(T data) {
        return new VDBOptResult<>(true, data);
    }

    private VDBOptResult<Void> fail() {
        return new VDBOptResult<>(false, null);
    }

    private <T> VDBOptResult<T> failData(T data) {
        return new VDBOptResult<>(false, data);
    }

    private <T> VDBOptResult<T> resultData(boolean success, T data) {
        return new VDBOptResult<>(success, data);
    }


    private String getUrl() {
        return getResource().getUrl();
    }

    private String getUsername() {
        return getResource().getUsername();
    }

    private String getPassowrd() {
        return getResource().getPassword();
    }

    @Override
    public VDBOptResult<Boolean> dropDatabase(String database) {
        String template = "DROP DATABASE IF EXISTS %s";
        String sql = String.format(template, database);
        Connection conn = null;
        try {
            conn = getConnFromResource();
            Boolean result = executeUpdate(conn, sql);
            return resultData(result, result);
        } catch (Exception e) {
            return new VDBOptResult<>(e.getMessage(), e);
        } finally {
            closeConnection(conn);
        }
    }
}
